

<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>dodo_detector.detection &mdash; dodo detector 0.6.1 documentation</title>
  

  
  
  
  

  
  <script type="text/javascript" src="../../_static/js/modernizr.min.js"></script>
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script type="text/javascript" src="../../_static/jquery.js"></script>
        <script type="text/javascript" src="../../_static/underscore.js"></script>
        <script type="text/javascript" src="../../_static/doctools.js"></script>
        <script type="text/javascript" src="../../_static/language_data.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    

  
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html" class="icon icon-home"> dodo detector
          

          
          </a>

          
            
            
              <div class="version">
                0.6.1
              </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">Contents:</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../modules.html">Package modules</a></li>
</ul>

            
          
        </div>
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../index.html">dodo detector</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> &raquo;</li>
        
          <li><a href="../index.html">Module code</a> &raquo;</li>
        
      <li>dodo_detector.detection</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for dodo_detector.detection</h1><div class="highlight"><pre>
<span></span><span class="ch">#!/usr/bin/env python</span>

<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">cv2</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">tensorflow</span> <span class="k">as</span> <span class="nn">tf</span>

<span class="kn">from</span> <span class="nn">tqdm</span> <span class="k">import</span> <span class="n">tqdm</span>
<span class="kn">from</span> <span class="nn">abc</span> <span class="k">import</span> <span class="n">ABCMeta</span><span class="p">,</span> <span class="n">abstractmethod</span>
<span class="kn">from</span> <span class="nn">object_detection.utils</span> <span class="k">import</span> <span class="n">label_map_util</span>
<span class="kn">from</span> <span class="nn">object_detection.utils</span> <span class="k">import</span> <span class="n">visualization_utils</span> <span class="k">as</span> <span class="n">vis_util</span>
<span class="kn">from</span> <span class="nn">imutils.video</span> <span class="k">import</span> <span class="n">WebcamVideoStream</span>
<span class="kn">from</span> <span class="nn">warnings</span> <span class="k">import</span> <span class="n">warn</span>


<div class="viewcode-block" id="ObjectDetector"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector">[docs]</a><span class="k">class</span> <span class="nc">ObjectDetector</span><span class="p">():</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Base class for object detectors used by the package.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">__metaclass__</span> <span class="o">=</span> <span class="n">ABCMeta</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="c1"># create logger</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="s1">&#39;dodo_detector&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_logger</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">logging</span><span class="o">.</span><span class="n">DEBUG</span><span class="p">)</span>
        <span class="c1"># create file handler which logs even debug messages</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_fh</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">FileHandler</span><span class="p">(</span><span class="s1">&#39;/tmp/dodo_detector.log&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_fh</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">logging</span><span class="o">.</span><span class="n">DEBUG</span><span class="p">)</span>
        <span class="c1"># create console handler with a higher log level</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_ch</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">StreamHandler</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_ch</span><span class="o">.</span><span class="n">setLevel</span><span class="p">(</span><span class="n">logging</span><span class="o">.</span><span class="n">DEBUG</span><span class="p">)</span>
        <span class="c1"># create formatter and add it to the handlers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_formatter</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">Formatter</span><span class="p">(</span><span class="s1">&#39;[</span><span class="si">%(asctime)s</span><span class="s1"> - </span><span class="si">%(name)s</span><span class="s1">]: </span><span class="si">%(levelname)s</span><span class="s1">: </span><span class="si">%(message)s</span><span class="s1">&#39;</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_fh</span><span class="o">.</span><span class="n">setFormatter</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_formatter</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_ch</span><span class="o">.</span><span class="n">setFormatter</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_formatter</span><span class="p">)</span>
        <span class="c1"># add the handlers to the logger</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_logger</span><span class="o">.</span><span class="n">addHandler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_fh</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_logger</span><span class="o">.</span><span class="n">addHandler</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_ch</span><span class="p">)</span>

<div class="viewcode-block" id="ObjectDetector.from_image"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector.from_image">[docs]</a>    <span class="nd">@abstractmethod</span>
    <span class="k">def</span> <span class="nf">from_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">frame</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Detects objects in an image</span>

<span class="sd">        :param frame: a numpy.ndarray containing the image where objects will be detected</span>
<span class="sd">        :return: a tuple containing the image, with objects marked by rectangles,</span>
<span class="sd">                 and a dictionary listing objects and their locations as `(ymin, xmin, ymax, xmax)`</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">pass</span></div>

    <span class="k">def</span> <span class="nf">_detect_from_stream</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">get_frame</span><span class="p">,</span> <span class="n">stream</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        This internal method detects objects from images retrieved from a stream, given a method that extracts frames from this stream</span>

<span class="sd">        :param get_frame: a method that extracts frames from the stream</span>
<span class="sd">        :param stream: an object representing a stream of images</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">ret</span><span class="p">,</span> <span class="n">frame</span> <span class="o">=</span> <span class="n">get_frame</span><span class="p">(</span><span class="n">stream</span><span class="p">)</span>

        <span class="k">while</span> <span class="n">ret</span><span class="p">:</span>
            <span class="n">marked_frame</span><span class="p">,</span> <span class="n">objects</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">from_image</span><span class="p">(</span><span class="n">frame</span><span class="p">)</span>

            <span class="n">cv2</span><span class="o">.</span><span class="n">imshow</span><span class="p">(</span><span class="s2">&quot;detection&quot;</span><span class="p">,</span> <span class="n">marked_frame</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">cv2</span><span class="o">.</span><span class="n">waitKey</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span> <span class="o">==</span> <span class="mi">27</span><span class="p">:</span>
                <span class="k">break</span>  <span class="c1"># ESC to quit</span>

            <span class="n">ret</span><span class="p">,</span> <span class="n">frame</span> <span class="o">=</span> <span class="n">get_frame</span><span class="p">(</span><span class="n">stream</span><span class="p">)</span>

        <span class="n">cv2</span><span class="o">.</span><span class="n">destroyAllWindows</span><span class="p">()</span>

<div class="viewcode-block" id="ObjectDetector.from_camera"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector.from_camera">[docs]</a>    <span class="k">def</span> <span class="nf">from_camera</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">camera_id</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Detects objects in frames from a camera feed</span>

<span class="sd">        :param camera_id: the ID of the camera in the system</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">def</span> <span class="nf">get_frame</span><span class="p">(</span><span class="n">stream</span><span class="p">):</span>
            <span class="n">frame</span> <span class="o">=</span> <span class="n">stream</span><span class="o">.</span><span class="n">read</span><span class="p">()</span>
            <span class="n">ret</span> <span class="o">=</span> <span class="kc">True</span>
            <span class="k">return</span> <span class="n">ret</span><span class="p">,</span> <span class="n">frame</span>

        <span class="n">stream</span> <span class="o">=</span> <span class="n">WebcamVideoStream</span><span class="p">(</span><span class="n">src</span><span class="o">=</span><span class="n">camera_id</span><span class="p">)</span>

        <span class="n">stream</span><span class="o">.</span><span class="n">start</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_detect_from_stream</span><span class="p">(</span><span class="n">get_frame</span><span class="p">,</span> <span class="n">stream</span><span class="p">)</span>
        <span class="n">stream</span><span class="o">.</span><span class="n">stop</span><span class="p">()</span></div>

<div class="viewcode-block" id="ObjectDetector.from_video"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector.from_video">[docs]</a>    <span class="k">def</span> <span class="nf">from_video</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">filepath</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Detects objects in frames from a video file</span>
<span class="sd">        </span>
<span class="sd">        :param filepath: the path to the video file</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">def</span> <span class="nf">get_frame</span><span class="p">(</span><span class="n">stream</span><span class="p">):</span>
            <span class="n">ret</span><span class="p">,</span> <span class="n">frame</span> <span class="o">=</span> <span class="n">stream</span><span class="o">.</span><span class="n">read</span><span class="p">()</span>
            <span class="k">return</span> <span class="n">ret</span><span class="p">,</span> <span class="n">frame</span>

        <span class="n">stream</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">VideoCapture</span><span class="p">()</span>
        <span class="n">stream</span><span class="o">.</span><span class="n">open</span><span class="p">(</span><span class="n">filename</span><span class="o">=</span><span class="n">filepath</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_detect_from_stream</span><span class="p">(</span><span class="n">get_frame</span><span class="p">,</span> <span class="n">stream</span><span class="p">)</span></div>

<div class="viewcode-block" id="ObjectDetector.is_rgb"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector.is_rgb">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">is_rgb</span><span class="p">(</span><span class="n">im</span><span class="p">):</span>
        <span class="k">return</span> <span class="nb">len</span><span class="p">(</span><span class="n">im</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">3</span> <span class="ow">and</span> <span class="n">im</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">==</span> <span class="mi">3</span></div>

<div class="viewcode-block" id="ObjectDetector.to_rgb"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.ObjectDetector.to_rgb">[docs]</a>    <span class="nd">@staticmethod</span>
    <span class="k">def</span> <span class="nf">to_rgb</span><span class="p">(</span><span class="n">im</span><span class="p">):</span>
        <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">im</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">im</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

        <span class="n">ret</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">empty</span><span class="p">((</span><span class="n">w</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">uint8</span><span class="p">)</span>
        <span class="n">ret</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="n">ret</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">ret</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">im</span>
        <span class="k">return</span> <span class="n">ret</span></div></div>


<div class="viewcode-block" id="KeypointObjectDetector"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.KeypointObjectDetector">[docs]</a><span class="k">class</span> <span class="nc">KeypointObjectDetector</span><span class="p">(</span><span class="n">ObjectDetector</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Object detector based on keypoints. This class depends on OpenCV SIFT and SURF feature detection algorithms,</span>
<span class="sd">    as well as the brute-force and FLANN-based feature matchers.</span>

<span class="sd">    :param database_path: Path to the top-level directory containing subdirectories, each subdirectory named after a class of objects and containing images of that object.</span>
<span class="sd">    :param detector_type: either `SURF`, `SIFT` or `RootSIFT`</span>
<span class="sd">    :param matcher_type: either `BF` for brute-force matcher or `FLANN` for flann-based matcher</span>
<span class="sd">    :param min_points: minimum number of keypoints necessary for an object to be considered detected in an image</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">detector_type</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detector_type</span>

    <span class="nd">@detector_type</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">detector_type</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="c1"># create the detector</span>
        <span class="k">if</span> <span class="n">value</span> <span class="ow">in</span> <span class="p">[</span><span class="s1">&#39;SIFT&#39;</span><span class="p">,</span> <span class="s1">&#39;RootSIFT&#39;</span><span class="p">]:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">detector</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">xfeatures2d</span><span class="o">.</span><span class="n">SIFT_create</span><span class="p">()</span>
        <span class="k">elif</span> <span class="n">value</span> <span class="o">==</span> <span class="s1">&#39;SURF&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">detector</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">xfeatures2d</span><span class="o">.</span><span class="n">SURF_create</span><span class="p">()</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Invalid detector type&#39;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_detector_type</span> <span class="o">=</span> <span class="n">value</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">matcher_type</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span>

    <span class="nd">@matcher_type</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">matcher_type</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="c1"># get which OpenCV feature matcher the user wants</span>
        <span class="k">if</span> <span class="n">value</span> <span class="o">==</span> <span class="s1">&#39;BF&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">matcher</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">BFMatcher</span><span class="p">()</span>
        <span class="k">elif</span> <span class="n">value</span> <span class="o">==</span> <span class="s1">&#39;FLANN&#39;</span><span class="p">:</span>
            <span class="n">flann_index_kdtree</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="n">index_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">algorithm</span><span class="o">=</span><span class="n">flann_index_kdtree</span><span class="p">,</span> <span class="n">trees</span><span class="o">=</span><span class="mi">5</span><span class="p">)</span>
            <span class="n">search_params</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">checks</span><span class="o">=</span><span class="mi">50</span><span class="p">)</span>  <span class="c1"># or pass empty dictionary</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">matcher</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">FlannBasedMatcher</span><span class="p">(</span><span class="n">index_params</span><span class="p">,</span> <span class="n">search_params</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Invalid matcher type&#39;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_matcher_type</span> <span class="o">=</span> <span class="n">value</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">database_path</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_path</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">categories</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_categories</span>

    <span class="nd">@database_path</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">database_path</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="c1"># get the directory where object textures are stored</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_database_path</span> <span class="o">=</span> <span class="n">value</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_path</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">!=</span> <span class="s1">&#39;/&#39;</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_database_path</span> <span class="o">+=</span> <span class="s1">&#39;/&#39;</span>
        <span class="c1"># store object classes in a list</span>
        <span class="c1"># each directory in the object database corresponds to a class</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_categories</span> <span class="o">=</span> <span class="p">[</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">basename</span><span class="p">(</span><span class="n">d</span><span class="p">)</span> <span class="k">for</span> <span class="n">d</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">listdir</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_database_path</span><span class="p">)]</span>
        <span class="c1"># minimum object dimensions in pixels</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_object_height</span> <span class="o">=</span> <span class="mi">10</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_object_width</span> <span class="o">=</span> <span class="mi">10</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_object_area</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_object_height</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_object_width</span>
        <span class="c1"># initialize the frame counter for each object class at 0</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">object_counters</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">ob</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">categories</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">object_counters</span><span class="p">[</span><span class="n">ob</span><span class="p">]</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="c1"># load features for each texture and store the image,</span>
        <span class="c1"># its keypoints and corresponding descriptor</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">object_features</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">obj</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">categories</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">object_features</span><span class="p">[</span><span class="n">obj</span><span class="p">]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_load_features</span><span class="p">(</span><span class="n">obj</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">database_path</span><span class="p">,</span> <span class="n">detector_type</span><span class="o">=</span><span class="s1">&#39;RootSIFT&#39;</span><span class="p">,</span> <span class="n">matcher_type</span><span class="o">=</span><span class="s1">&#39;BF&#39;</span><span class="p">,</span> <span class="n">min_points</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">logging</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ObjectDetector</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">current_frame</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="c1"># these things are properties, take a look at their setters in this class</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">detector_type</span> <span class="o">=</span> <span class="n">detector_type</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">matcher_type</span> <span class="o">=</span> <span class="n">matcher_type</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">database_path</span> <span class="o">=</span> <span class="n">database_path</span>

        <span class="c1"># minimum number of features for a KNN match to consider that an object has been found</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">min_points</span> <span class="o">=</span> <span class="n">min_points</span>

    <span class="k">def</span> <span class="nf">_load_features</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">object_name</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Given the name of an object class from the image database directory, this method iterates through all the images contained in that directory and extracts their keypoints and descriptors using the desired feature detector</span>

<span class="sd">        :param object_name: the name of an object class, whose image directory is contained inside the image database directory</span>
<span class="sd">        :return: a list of tuples, each tuple containing the processed image as a grayscale numpy.ndarray, its keypoints and desciptors</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">img_files</span> <span class="o">=</span> <span class="p">[</span>
            <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">database_path</span> <span class="o">+</span> <span class="n">object_name</span> <span class="o">+</span> <span class="s1">&#39;/&#39;</span><span class="p">,</span> <span class="n">f</span><span class="p">)</span> <span class="k">for</span> <span class="n">f</span> <span class="ow">in</span> <span class="n">os</span><span class="o">.</span><span class="n">listdir</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">database_path</span> <span class="o">+</span> <span class="n">object_name</span> <span class="o">+</span> <span class="s1">&#39;/&#39;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">isfile</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">database_path</span> <span class="o">+</span> <span class="n">object_name</span> <span class="o">+</span> <span class="s1">&#39;/&#39;</span><span class="p">,</span> <span class="n">f</span><span class="p">))</span>
        <span class="p">]</span>

        <span class="n">pbar</span> <span class="o">=</span> <span class="n">tqdm</span><span class="p">(</span><span class="n">desc</span><span class="o">=</span><span class="n">object_name</span><span class="p">,</span> <span class="n">total</span><span class="o">=</span><span class="nb">len</span><span class="p">(</span><span class="n">img_files</span><span class="p">))</span>

        <span class="c1"># extract the keypoints from all images in the database</span>
        <span class="n">features</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">img_file</span> <span class="ow">in</span> <span class="n">img_files</span><span class="p">:</span>
            <span class="n">pbar</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
            <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">imread</span><span class="p">(</span><span class="n">img_file</span><span class="p">)</span>

            <span class="c1"># scaling_factor = 640 / img.shape[0]</span>
            <span class="k">if</span> <span class="n">img</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">&gt;</span> <span class="mi">1000</span><span class="p">:</span>
                <span class="n">img</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">resize</span><span class="p">(</span><span class="n">img</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="n">fx</span><span class="o">=</span><span class="mf">0.3</span><span class="p">,</span> <span class="n">fy</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>

            <span class="c1"># find keypoints and descriptors with the selected feature detector</span>
            <span class="n">keypoints</span><span class="p">,</span> <span class="n">descriptors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detectAndCompute</span><span class="p">(</span><span class="n">img</span><span class="p">)</span>

            <span class="n">features</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">img</span><span class="p">,</span> <span class="n">keypoints</span><span class="p">,</span> <span class="n">descriptors</span><span class="p">))</span>

        <span class="k">return</span> <span class="n">features</span>

    <span class="k">def</span> <span class="nf">_detect_object</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">,</span> <span class="n">img_features</span><span class="p">,</span> <span class="n">scene</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>

<span class="sd">        :param name: name of the object class</span>
<span class="sd">        :param img_features: a list of tuples, each tuple containing three elements: an image, its keypoints and its descriptors.</span>
<span class="sd">        :param scene: the image where the object `name` will be detected</span>
<span class="sd">        :return: a tuple containing two elements: the homography matrix and the coordinates of a rectangle containing the object in a list `[xmin, ymin, xmax, ymax]`</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">scene_kp</span><span class="p">,</span> <span class="n">scene_descs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detectAndCompute</span><span class="p">(</span><span class="n">scene</span><span class="p">)</span>

        <span class="k">for</span> <span class="n">img_feature</span> <span class="ow">in</span> <span class="n">img_features</span><span class="p">:</span>
            <span class="n">obj_image</span><span class="p">,</span> <span class="n">obj_keypoints</span><span class="p">,</span> <span class="n">obj_descriptors</span> <span class="o">=</span> <span class="n">img_feature</span>

            <span class="k">if</span> <span class="n">obj_descriptors</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">obj_descriptors</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">and</span> <span class="n">scene_descs</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">scene_descs</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">matches</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">matcher</span><span class="o">.</span><span class="n">knnMatch</span><span class="p">(</span><span class="n">obj_descriptors</span><span class="p">,</span> <span class="n">scene_descs</span><span class="p">,</span> <span class="n">k</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>

                <span class="c1"># store all the good matches as per Lowe&#39;s ratio test</span>
                <span class="n">good</span> <span class="o">=</span> <span class="p">[]</span>
                <span class="k">for</span> <span class="n">match</span> <span class="ow">in</span> <span class="n">matches</span><span class="p">:</span>
                    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">match</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
                        <span class="n">m</span><span class="p">,</span> <span class="n">n</span> <span class="o">=</span> <span class="n">match</span>
                        <span class="k">if</span> <span class="n">m</span><span class="o">.</span><span class="n">distance</span> <span class="o">&lt;</span> <span class="mf">0.7</span> <span class="o">*</span> <span class="n">n</span><span class="o">.</span><span class="n">distance</span><span class="p">:</span>
                            <span class="n">good</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">m</span><span class="p">)</span>

                <span class="c1"># an object was detected</span>
                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">good</span><span class="p">)</span> <span class="o">&gt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">min_points</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">object_counters</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">+=</span> <span class="mi">1</span>
                    <span class="n">source_points</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">([</span><span class="n">obj_keypoints</span><span class="p">[</span><span class="n">m</span><span class="o">.</span><span class="n">queryIdx</span><span class="p">]</span><span class="o">.</span><span class="n">pt</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">good</span><span class="p">])</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
                    <span class="n">destination_points</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">([</span><span class="n">scene_kp</span><span class="p">[</span><span class="n">m</span><span class="o">.</span><span class="n">trainIdx</span><span class="p">]</span><span class="o">.</span><span class="n">pt</span> <span class="k">for</span> <span class="n">m</span> <span class="ow">in</span> <span class="n">good</span><span class="p">])</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>

                    <span class="n">M</span><span class="p">,</span> <span class="n">mask</span> <span class="o">=</span> <span class="n">cv2</span><span class="o">.</span><span class="n">findHomography</span><span class="p">(</span><span class="n">source_points</span><span class="p">,</span> <span class="n">destination_points</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">RANSAC</span><span class="p">,</span> <span class="mf">5.0</span><span class="p">)</span>

                    <span class="k">if</span> <span class="n">M</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                        <span class="k">break</span>

                    <span class="k">if</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">obj_image</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">):</span>
                        <span class="n">h</span><span class="p">,</span> <span class="n">w</span> <span class="o">=</span> <span class="n">obj_image</span><span class="o">.</span><span class="n">shape</span>
                    <span class="k">else</span><span class="p">:</span>
                        <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">obj_image</span><span class="o">.</span><span class="n">shape</span>

                    <span class="n">pts</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="n">w</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="n">w</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">]])</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>

                    <span class="c1"># dst contains the coordinates of the vertices of the drawn rectangle</span>
                    <span class="n">dst</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">(</span><span class="n">cv2</span><span class="o">.</span><span class="n">perspectiveTransform</span><span class="p">(</span><span class="n">pts</span><span class="p">,</span> <span class="n">M</span><span class="p">))</span>

                    <span class="c1"># get the min and max x and y coordinates of the object</span>
                    <span class="n">xmin</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">dst</span><span class="p">[</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">4</span><span class="p">))</span>
                    <span class="n">xmax</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">dst</span><span class="p">[</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">4</span><span class="p">))</span>
                    <span class="n">ymin</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">dst</span><span class="p">[</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">4</span><span class="p">))</span>
                    <span class="n">ymax</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">dst</span><span class="p">[</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">][</span><span class="mi">1</span><span class="p">]</span> <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">4</span><span class="p">))</span>

                    <span class="c1"># transform homography into a simpler data structure</span>
                    <span class="n">dst</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="n">point</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">point</span> <span class="ow">in</span> <span class="n">dst</span><span class="p">],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">)</span>

                    <span class="c1"># returns the homography and a rectangle containing the object</span>
                    <span class="k">return</span> <span class="n">dst</span><span class="p">,</span> <span class="p">[</span><span class="n">xmin</span><span class="p">,</span> <span class="n">ymin</span><span class="p">,</span> <span class="n">xmax</span><span class="p">,</span> <span class="n">ymax</span><span class="p">]</span>

        <span class="k">return</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span>

    <span class="k">def</span> <span class="nf">_detectAndCompute</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">image</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">        Detects keypoints and generates descriptors according to the desired algorithm</span>
<span class="sd">        </span>
<span class="sd">        :param image: a numpy.ndarray containing the image whose keypoints and descriptors will be processed</span>
<span class="sd">        :return: a tuple containing keypoints and descriptors</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">keypoints</span><span class="p">,</span> <span class="n">descriptors</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">detector</span><span class="o">.</span><span class="n">detectAndCompute</span><span class="p">(</span><span class="n">image</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">detector_type</span> <span class="o">==</span> <span class="s1">&#39;RootSIFT&#39;</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">keypoints</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="c1"># Transforms SIFT descriptors into RootSIFT descriptors</span>
            <span class="c1"># apply the Hellinger kernel by first L1-normalizing</span>
            <span class="c1"># and taking the square-root</span>
            <span class="n">eps</span> <span class="o">=</span> <span class="mf">1e-7</span>
            <span class="n">descriptors</span> <span class="o">/=</span> <span class="p">(</span><span class="n">descriptors</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">keepdims</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="o">+</span> <span class="n">eps</span><span class="p">)</span>
            <span class="n">descriptors</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">descriptors</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">keypoints</span><span class="p">,</span> <span class="n">descriptors</span>

<div class="viewcode-block" id="KeypointObjectDetector.from_image"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.KeypointObjectDetector.from_image">[docs]</a>    <span class="k">def</span> <span class="nf">from_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">frame</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">current_frame</span> <span class="o">+=</span> <span class="mi">1</span>
        <span class="c1"># Our operations on the frame come here</span>

        <span class="n">detected_objects</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="k">for</span> <span class="n">object_name</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">object_features</span><span class="p">:</span>
            <span class="n">homography</span><span class="p">,</span> <span class="n">rct</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detect_object</span><span class="p">(</span><span class="n">object_name</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">object_features</span><span class="p">[</span><span class="n">object_name</span><span class="p">],</span> <span class="n">frame</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">rct</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">xmin</span> <span class="o">=</span> <span class="n">rct</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">ymin</span> <span class="o">=</span> <span class="n">rct</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="n">xmax</span> <span class="o">=</span> <span class="n">rct</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
                <span class="n">ymax</span> <span class="o">=</span> <span class="n">rct</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>

                <span class="k">if</span> <span class="n">object_name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">detected_objects</span><span class="p">:</span>
                    <span class="n">detected_objects</span><span class="p">[</span><span class="n">object_name</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>

                <span class="n">detected_objects</span><span class="p">[</span><span class="n">object_name</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">({</span><span class="s1">&#39;box&#39;</span><span class="p">:</span> <span class="p">(</span><span class="n">ymin</span><span class="p">,</span> <span class="n">xmin</span><span class="p">,</span> <span class="n">ymax</span><span class="p">,</span> <span class="n">xmax</span><span class="p">)})</span>

                <span class="n">text_point</span> <span class="o">=</span> <span class="p">(</span><span class="n">homography</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">homography</span><span class="p">[</span><span class="mi">1</span><span class="p">][</span><span class="mi">1</span><span class="p">])</span>
                <span class="n">homography</span> <span class="o">=</span> <span class="n">homography</span><span class="o">.</span><span class="n">reshape</span><span class="p">((</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
                <span class="n">cv2</span><span class="o">.</span><span class="n">polylines</span><span class="p">(</span><span class="n">frame</span><span class="p">,</span> <span class="p">[</span><span class="n">homography</span><span class="p">],</span> <span class="kc">True</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">255</span><span class="p">,</span> <span class="mi">255</span><span class="p">),</span> <span class="mi">10</span><span class="p">)</span>

                <span class="n">cv2</span><span class="o">.</span><span class="n">putText</span><span class="p">(</span><span class="n">frame</span><span class="p">,</span> <span class="n">object_name</span> <span class="o">+</span> <span class="s1">&#39;: &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">object_counters</span><span class="p">[</span><span class="n">object_name</span><span class="p">]),</span> <span class="n">text_point</span><span class="p">,</span> <span class="n">cv2</span><span class="o">.</span><span class="n">FONT_HERSHEY_COMPLEX_SMALL</span><span class="p">,</span> <span class="mf">1.2</span><span class="p">,</span> <span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="mi">2</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">frame</span><span class="p">,</span> <span class="n">detected_objects</span></div></div>


<div class="viewcode-block" id="SingleShotDetector"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.SingleShotDetector">[docs]</a><span class="k">class</span> <span class="nc">SingleShotDetector</span><span class="p">(</span><span class="n">ObjectDetector</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Object detector powered by the TensorFlow Object Detection API.</span>

<span class="sd">    :param path_to_frozen_graph: path to the frozen inference graph file, a file with a `.pb` extension.</span>
<span class="sd">    :param path_to_labels: path to the label map, a text file with the `.pbtxt` extension.</span>
<span class="sd">    :param num_classes: number of object classes that will be detected. If None, it will be guessed by the contents of the label map.</span>
<span class="sd">    :param confidence: a value between 0 and 1 representing the confidence level the network has in the detection to consider it an actual detection.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">path_to_frozen_graph</span><span class="p">,</span> <span class="n">path_to_labels</span><span class="p">,</span> <span class="n">num_classes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">confidence</span><span class="o">=.</span><span class="mi">8</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ObjectDetector</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="mi">0</span> <span class="o">&lt;</span> <span class="n">confidence</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;confidence must be between 0 and 1&quot;</span><span class="p">)</span>

        <span class="c1"># load (frozen) tensorflow model into memory</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">Graph</span><span class="p">()</span>
        <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">as_default</span><span class="p">():</span>
            <span class="n">od_graph_def</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">GraphDef</span><span class="p">()</span>
            <span class="k">with</span> <span class="n">tf</span><span class="o">.</span><span class="n">gfile</span><span class="o">.</span><span class="n">GFile</span><span class="p">(</span><span class="n">path_to_frozen_graph</span><span class="p">,</span> <span class="s1">&#39;rb&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">fid</span><span class="p">:</span>
                <span class="n">serialized_graph</span> <span class="o">=</span> <span class="n">fid</span><span class="o">.</span><span class="n">read</span><span class="p">()</span>
                <span class="n">od_graph_def</span><span class="o">.</span><span class="n">ParseFromString</span><span class="p">(</span><span class="n">serialized_graph</span><span class="p">)</span>
                <span class="n">tf</span><span class="o">.</span><span class="n">import_graph_def</span><span class="p">(</span><span class="n">od_graph_def</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s1">&#39;&#39;</span><span class="p">)</span>

        <span class="c1"># Label maps map indices to category names, so that when our convolution</span>
        <span class="c1"># network predicts 5, we know that this corresponds to airplane.</span>
        <span class="c1"># Here we use internal utility functions, but anything that returns a</span>
        <span class="c1"># dictionary mapping integers to appropriate string labels would be fine</span>
        <span class="n">label_map</span> <span class="o">=</span> <span class="n">label_map_util</span><span class="o">.</span><span class="n">load_labelmap</span><span class="p">(</span><span class="n">path_to_labels</span><span class="p">)</span>

        <span class="c1"># this is a workaround to guess the number of classes by the contents of the label map</span>
        <span class="c1"># it may not be perfect</span>
        <span class="n">label_map_contents</span> <span class="o">=</span> <span class="nb">open</span><span class="p">(</span><span class="n">path_to_labels</span><span class="p">,</span> <span class="s1">&#39;r&#39;</span><span class="p">)</span><span class="o">.</span><span class="n">read</span><span class="p">()</span>
        <span class="n">suggested_num_classes</span> <span class="o">=</span> <span class="n">label_map_contents</span><span class="o">.</span><span class="n">count</span><span class="p">(</span><span class="s1">&#39;name:&#39;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">num_classes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">num_classes</span> <span class="o">=</span> <span class="n">suggested_num_classes</span>
        <span class="k">elif</span> <span class="n">num_classes</span> <span class="o">!=</span> <span class="n">suggested_num_classes</span><span class="p">:</span>
            <span class="c1"># a warning is generated in case the user passes a number of classes</span>
            <span class="c1"># that is different than the label map contents</span>
            <span class="n">warn</span><span class="p">(</span><span class="s1">&#39;Suggested number of classes according to label map is </span><span class="si">{0}</span><span class="s1">. Will use </span><span class="si">{1}</span><span class="s1"> as manually passed.&#39;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">suggested_num_classes</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">))</span>

        <span class="n">categories</span> <span class="o">=</span> <span class="n">label_map_util</span><span class="o">.</span><span class="n">convert_label_map_to_categories</span><span class="p">(</span><span class="n">label_map</span><span class="p">,</span> <span class="n">max_num_classes</span><span class="o">=</span><span class="n">num_classes</span><span class="p">,</span> <span class="n">use_display_name</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_category_index</span> <span class="o">=</span> <span class="n">label_map_util</span><span class="o">.</span><span class="n">create_category_index</span><span class="p">(</span><span class="n">categories</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_categories</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_categories_public</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">tmp</span> <span class="ow">in</span> <span class="n">categories</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_categories</span><span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">tmp</span><span class="p">[</span><span class="s1">&#39;id&#39;</span><span class="p">])]</span> <span class="o">=</span> <span class="n">tmp</span><span class="p">[</span><span class="s1">&#39;name&#39;</span><span class="p">]</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_categories_public</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">tmp</span><span class="p">[</span><span class="s1">&#39;name&#39;</span><span class="p">])</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">_confidence</span> <span class="o">=</span> <span class="n">confidence</span>

        <span class="c1"># create a session that will be used until our detector is set on fire by the gc</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_session</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">Session</span><span class="p">(</span><span class="n">graph</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="p">)</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">confidence</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_confidence</span>

    <span class="nd">@property</span>
    <span class="k">def</span> <span class="nf">categories</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_categories_public</span>

    <span class="nd">@confidence</span><span class="o">.</span><span class="n">setter</span>
    <span class="k">def</span> <span class="nf">confidence</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">value</span><span class="p">):</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">_confidence</span> <span class="o">=</span> <span class="n">value</span>

<div class="viewcode-block" id="SingleShotDetector.from_image"><a class="viewcode-back" href="../../dodo_detector.html#dodo_detector.detection.SingleShotDetector.from_image">[docs]</a>    <span class="k">def</span> <span class="nf">from_image</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">frame</span><span class="p">):</span>
        <span class="c1"># object recognition begins here</span>

        <span class="k">if</span> <span class="ow">not</span> <span class="n">ObjectDetector</span><span class="o">.</span><span class="n">is_rgb</span><span class="p">(</span><span class="n">frame</span><span class="p">):</span>
            <span class="n">frame</span> <span class="o">=</span> <span class="n">ObjectDetector</span><span class="o">.</span><span class="n">to_rgb</span><span class="p">(</span><span class="n">frame</span><span class="p">)</span>

        <span class="n">height</span><span class="p">,</span> <span class="n">width</span><span class="p">,</span> <span class="n">z</span> <span class="o">=</span> <span class="n">frame</span><span class="o">.</span><span class="n">shape</span>

        <span class="n">image_np_expanded</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">frame</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
        <span class="n">image_tensor</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">get_tensor_by_name</span><span class="p">(</span><span class="s1">&#39;image_tensor:0&#39;</span><span class="p">)</span>
        <span class="c1"># Each box represents a part of the image where a particular object was detected.</span>
        <span class="n">boxes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">get_tensor_by_name</span><span class="p">(</span><span class="s1">&#39;detection_boxes:0&#39;</span><span class="p">)</span>
        <span class="c1"># Each score represent how level of confidence for each of the objects.</span>
        <span class="c1"># Score is shown on the result image, together with the class label.</span>
        <span class="n">scores</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">get_tensor_by_name</span><span class="p">(</span><span class="s1">&#39;detection_scores:0&#39;</span><span class="p">)</span>
        <span class="n">classes</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">get_tensor_by_name</span><span class="p">(</span><span class="s1">&#39;detection_classes:0&#39;</span><span class="p">)</span>
        <span class="n">num_detections</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_detection_graph</span><span class="o">.</span><span class="n">get_tensor_by_name</span><span class="p">(</span><span class="s1">&#39;num_detections:0&#39;</span><span class="p">)</span>

        <span class="c1"># Actual detection</span>
        <span class="n">boxes</span><span class="p">,</span> <span class="n">scores</span><span class="p">,</span> <span class="n">classes</span><span class="p">,</span> <span class="n">num_detections</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_session</span><span class="o">.</span><span class="n">run</span><span class="p">([</span><span class="n">boxes</span><span class="p">,</span> <span class="n">scores</span><span class="p">,</span> <span class="n">classes</span><span class="p">,</span> <span class="n">num_detections</span><span class="p">],</span> <span class="n">feed_dict</span><span class="o">=</span><span class="p">{</span><span class="n">image_tensor</span><span class="p">:</span> <span class="n">image_np_expanded</span><span class="p">})</span>

        <span class="c1"># count how many scores are above the designated threshold</span>
        <span class="n">worthy_detections</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">score</span> <span class="o">&gt;=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_confidence</span> <span class="k">for</span> <span class="n">score</span> <span class="ow">in</span> <span class="n">scores</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
        <span class="c1"># self._logger.debug(&#39;Found &#39; + str(worthy_detections) + &#39; objects&#39;)</span>

        <span class="n">detected_objects</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="c1"># analyze all worthy detections</span>
        <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">worthy_detections</span><span class="p">):</span>

            <span class="c1"># capture the class of the detected object</span>
            <span class="n">class_name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_categories</span><span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="n">classes</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">x</span><span class="p">])]</span>

            <span class="c1"># get the detection box around the object</span>
            <span class="n">box_objects</span> <span class="o">=</span> <span class="n">boxes</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">x</span><span class="p">]</span>

            <span class="c1"># positions of the box are between 0 and 1, relative to the size of the image</span>
            <span class="c1"># we multiply them by the size of the image to get the box location in pixels</span>
            <span class="n">ymin</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">box_objects</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">*</span> <span class="n">height</span><span class="p">)</span>
            <span class="n">xmin</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">box_objects</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span> <span class="o">*</span> <span class="n">width</span><span class="p">)</span>
            <span class="n">ymax</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">box_objects</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span> <span class="o">*</span> <span class="n">height</span><span class="p">)</span>
            <span class="n">xmax</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">box_objects</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span> <span class="o">*</span> <span class="n">width</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">class_name</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">detected_objects</span><span class="p">:</span>
                <span class="n">detected_objects</span><span class="p">[</span><span class="n">class_name</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>

            <span class="n">detected_objects</span><span class="p">[</span><span class="n">class_name</span><span class="p">]</span><span class="o">.</span><span class="n">append</span><span class="p">({</span><span class="s1">&#39;box&#39;</span><span class="p">:</span> <span class="p">(</span><span class="n">ymin</span><span class="p">,</span> <span class="n">xmin</span><span class="p">,</span> <span class="n">ymax</span><span class="p">,</span> <span class="n">xmax</span><span class="p">),</span> <span class="s1">&#39;confidence&#39;</span><span class="p">:</span> <span class="n">scores</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="n">x</span><span class="p">]})</span>

        <span class="c1"># Visualization of the results of a detection.</span>
        <span class="n">vis_util</span><span class="o">.</span><span class="n">visualize_boxes_and_labels_on_image_array</span><span class="p">(</span>
            <span class="n">frame</span><span class="p">,</span>
            <span class="n">np</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="n">boxes</span><span class="p">),</span>
            <span class="n">np</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="n">classes</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">),</span>
            <span class="n">np</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="n">scores</span><span class="p">),</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">_category_index</span><span class="p">,</span>
            <span class="n">use_normalized_coordinates</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
            <span class="n">line_thickness</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span>
            <span class="n">min_score_thresh</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_confidence</span>
        <span class="p">)</span>

        <span class="k">return</span> <span class="n">frame</span><span class="p">,</span> <span class="n">detected_objects</span></div></div>
</pre></div>

           </div>
           
          </div>
          <footer>
  

  <hr/>

  <div role="contentinfo">
    <p>
        &copy; Copyright 2019, Douglas De Rizzo Meneghetti

    </p>
  </div>
  Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>

        </div>
      </div>

    </section>

  </div>
  


  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>